import os
import sys
import threading
from docx import Document
import xlrd
import pptx
import PyPDF2
import odf.opendocument
from odf.text import P
from odf.table import TableRow, TableCell
import rtfparse
import PySimpleGUI as sg

def search_documents(keywords, root_dir, log_callback=None, result_callback=None):
    """搜索指定目录下的文档，查找包含关键词的文件"""
    if log_callback:
        log_callback(f"开始搜索目录: {root_dir}")
    
    result = []
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            file_extension = os.path.splitext(file)[1]
            file_path = os.path.join(root, file)
            
            if log_callback:
                log_callback(f"正在检查: {file_path}")
            
            try:
                content = ""
                if file_extension in ['.doc', '.docx']:
                    doc = Document(file_path)
                    content = ''.join([para.text for para in doc.paragraphs])
                elif file_extension in ['.xls', '.xlsx', '.ods']:
                    if file_extension in ['.xls', '.xlsx']:
                        workbook = xlrd.open_workbook(file_path)
                        for sheet in workbook.sheets():
                            for row in range(sheet.nrows):
                                content += ''.join(str(cell.value) for cell in sheet.row(row))
                    elif file_extension == '.ods':
                        doc = odf.opendocument.load(file_path)
                        for table in doc.spreadsheet.getElementsByType(odf.table.Table):
                            for row in table.getElementsByType(TableRow):
                                for cell in row.getElementsByType(TableCell):
                                    for p in cell.getElementsByType(P):
                                        content += ''.join(p.data)
                elif file_extension in ['.ppt', '.pptx', '.odp']:
                    if file_extension in ['.ppt', '.pptx']:
                        prs = pptx.Presentation(file_path)
                        content = ''.join([shape.text_frame.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, 'text_frame')])
                    elif file_extension == '.odp':
                        doc = odf.opendocument.load(file_path)
                        for p in doc.getElementsByType(P):
                            content += ''.join(p.data)
                elif file_extension == '.pdf':
                    with open(file_path, 'rb') as pdf_file:
                        pdf_reader = PyPDF2.PdfReader(pdf_file)
                        content = ''.join([page.extract_text() for page in pdf_reader.pages])
                elif file_extension == '.txt':
                    with open(file_path, 'r', encoding='utf-8', errors='ignore') as txt_file:
                        content = txt_file.read()
                elif file_extension == '.rtf':
                    with open(file_path, 'rb') as rtf_file:
                        rtf = rtfparse.from_file(rtf_file)
                        content = rtf.get_text()
                elif file_extension == '.csv':
                    with open(file_path, 'r', encoding='utf-8', errors='ignore') as csv_file:
                        content = csv_file.read()
                
                # 检查是否包含关键词
                if any(keyword in content for keyword in keywords):
                    result.append((file_path, content))
                    if result_callback:
                        result_callback(file_path, content, keywords)
            except Exception as e:
                if log_callback:
                    log_callback(f"读取文件 {file_path} 时出错: {e}")
    
    if log_callback:
        log_callback(f"搜索完成，共找到 {len(result)} 个匹配文件")
    
    return result

def get_context(text, keyword, before=2, after=2):
    """获取关键词前后的上下文"""
    sentences = text.replace('\n', ' ').split('。')
    context = []
    
    for i, sentence in enumerate(sentences):
        if keyword in sentence:
            start = max(0, i - before)
            end = min(len(sentences), i + after + 1)
            context.append('。'.join(sentences[start:end]) + '。')
    
    return context

def worker(keywords, root_dir, log_callback, result_callback, completion_callback):
    """后台工作线程，执行搜索任务"""
    try:
        results = search_documents(keywords, root_dir, log_callback, result_callback)
        if completion_callback:
            completion_callback(results)
    except Exception as e:
        if log_callback:
            log_callback(f"搜索过程中发生错误: {e}")

def main():
    """主函数，创建GUI界面"""
    # 设置中文字体支持
    sg.set_options(font=('SimHei', 10))
    
    # 布局定义
    layout = [
        [
            sg.Text("关键字：", size=(8, 1)),
            sg.InputText("秘密, 机密, 绝密", size=(50, 1), key='keywords'),
            sg.Button("开始", size=(10, 1), key='start')
        ],
        [
            sg.Multiline(size=(60, 30), font=('SimHei', 10), key='log', autoscroll=True, 
                         reroute_stdout=True, reroute_stderr=True, write_only=True),
            sg.Multiline(size=(60, 30), font=('SimHei', 10), key='results', autoscroll=True, 
                         write_only=True)
        ],
        [sg.Text("文山 张泽山@2025", justification='center', expand_x=True)]
    ]
    
    # 创建窗口
    window = sg.Window("文档关键词搜索工具", layout, finalize=True)
    
    # 获取控件引用
    log_output = window['log']
    results_output = window['results']
    start_button = window['start']
    
    def update_log(message):
        """更新日志输出"""
        log_output.update(message + '\n', append=True)
    
    def update_results(file_path, content, keywords):
        """更新搜索结果输出"""
        results_output.update(f"找到匹配文件: {file_path}\n", append=True)
        for keyword in keywords:
            if keyword in content:
                context = get_context(content, keyword)
                for c in context:
                    results_output.update(f"  关键词: '{keyword}'\n", append=True)
                    results_output.update(f"  上下文: {c}\n\n", append=True)
    
    def search_complete(results):
        """搜索完成回调"""
        update_log(f"搜索完成，共找到 {len(results)} 个匹配文件")
        start_button.update(disabled=False)
    
    # 主事件循环
    while True:
        event, values = window.read()
        
        if event == sg.WINDOW_CLOSED:
            break
        
        if event == 'start':
            # 禁用按钮防止重复点击
            start_button.update(disabled=True)
            
            # 清空之前的结果
            log_output.update('')
            results_output.update('')
            
            # 获取用户输入的关键词
            keywords_text = values['keywords']
            # 处理用户输入的关键词，去除引号和空格
            keywords = [kw.strip().replace("'", "").replace('"', '') for kw in keywords_text.split(',') if kw.strip()]
            
            if not keywords:
                update_log("请输入有效的关键词")
                start_button.update(disabled=False)
                continue
            
            update_log(f"开始搜索关键词: {', '.join(keywords)}")
            
            # 获取所有磁盘/挂载点
            if os.name == 'nt':  # Windows 系统
                import string
                disks = [f"{d}:\\" for d in string.ascii_uppercase if os.path.exists(f"{d}:\\")]
            elif os.name == 'posix':  # Linux 系统
                disks = []
                try:
                    with open('/proc/mounts', 'r') as f:
                        for line in f:
                            fields = line.strip().split()
                            if len(fields) >= 2:
                                mount_point = fields[1]
                                if os.path.ismount(mount_point):
                                    disks.append(mount_point)
                    update_log(f"找到 {len(disks)} 个挂载点")
                except Exception as e:
                    update_log(f"读取挂载点信息时出错: {e}")
                    disks = ['/']  # 默认从根目录开始搜索
            else:
                update_log("不支持的操作系统")
                start_button.update(disabled=False)
                continue
            
            # 在后台线程中执行搜索
            threading.Thread(
                target=worker,
                args=(keywords, '/', update_log, update_results, search_complete),
                daemon=True
            ).start()
    
    window.close()

if __name__ == "__main__":
    main()    